from mxnet.gluon import nn
from mxnet import nd


def mlpconv(channels, kernel_size, padding, strides=1, max_pooling=True):
    out = nn.Sequential()
    out.add(
        nn.Conv2D(channels=channels, kernel_size=kernel_size,
                  strides=strides, padding=padding,
                  activation='relu'),
        nn.Conv2D(channels=channels, kernel_size=1,
                  padding=0, strides=1, activation='relu'),
        nn.Conv2D(channels=channels, kernel_size=1,
                  padding=0, strides=1, activation='relu')
    )
    if max_pooling:
        out.add(nn.MaxPool2D(pool_size=3, strides=2))
    return out


blk = mlpconv(64, 3, 0)
blk.initialize()

x = nd.random.uniform(shape=(32, 3, 16, 16))
y = blk(x)
print(y.shape)

# conv to last, then avg pooling
net = nn.Sequential()
with net.name_scope():
    net.add(
        mlpconv(96, 11, 0, strides=4),
        mlpconv(256, 5, 2),
        mlpconv(384, 3, 1),
        nn.Dropout(.5),
        # 10 classes
        mlpconv(10, 3, 1, max_pooling=False),
        nn.AvgPool2D(pool_size=5),
        nn.Flatten()
    )

